from functools import reduce
import math
import tensorflow as tf
import tensorflow.keras as keras
import tensorflow.keras.backend as K
from tensorflow.keras import layers
from common import (Focus, Conv, C3, Bottleneck, SPP, Detect,
                    DWConv, C3TR, BottleneckCSP, TransformerLayer, TransformerBlock)

def yolov5_backbone(x,gw=2,gd=1):                #! 以yolov5s为例
    y1  = Focus(x,       3,  16*gw, 3,           node=0)       #   1 (160,160, 32)
    y2  = Conv (     16*gw,  32*gw, 3,        2, node=1)(y1)   #   2 (160,160, 64)
    y3  = C3   (y2,  32*gw,  32*gw, 1*gd,        node=2)       #   3 (160,160, 64)
    y4  = Conv (     32*gw,  64*gw, 3,        2, node=3)(y3)   #   4 ( 80, 80,128)
    y5  = C3   (y4,  64*gw,  64*gw, 3*gd,        node=4)       #!  5 ( 80, 80,128) P3
    y6  = Conv (     64*gw, 128*gw, 3,        2, node=5)(y5)   #   6 ( 40, 40,256)
    y7  = C3   (y6, 128*gw, 128*gw, 3*gd,        node=6)       #!  7 ( 40, 40,256) P4
    y8  = Conv (    128*gw, 256*gw, 3,        2, node=7)(y7)   #   8 ( 20, 20,512)
    y9  = SPP  (y8, 256*gw, 256*gw,              node=8)       #   9 ( 20, 20,512) 
    y10 = C3   (y9, 256*gw, 256*gw, 1*gd, False, node=9)       #! 10 ( 20, 20,512) P5
    return y5,y7,y10

def yolov5_head(y5,y7,y10,nc=80,imgsz=(640,640),strides=[8,16,32],anchors=None,ch=[128,256,512],gw=2,gd=1):
    y11 = Conv (    256*gw, 128*gw, 1,        1, node=10)(y10) #   11 (20,20,256) P5
    y12 = layers.UpSampling2D(name="11.UpSample")(y11)         #!  12 (40,40,256) P5_Upsample
    y13 = layers.Concatenate(name='12.Concat')([y12,y7])       #   13 (40,40,512) P5_Upsample
    y14 = C3   (y13,256*gw, 128*gw, 1*gd, False, node=13)      #   14 (40,40,256) P5_Upsample

    y15 = Conv (    128*gw,  64*gw, 1,        1, node=14)(y14) #   15 (40,40,128) P4
    y16 = layers.UpSampling2D(name='15.UpSample')(y15)         #!  16 (80,80,128) P4_Upsample
    y17 = layers.Concatenate(name='16.Concat')([y16,y5])       #   17 (80,80,256) P4_Upsample
    y18 = C3   (y17,128*gw,  64*gw, 1*gd, False, node=17)      #!  18 (80,80,128) P3_out

    y19 = Conv (     64*gw,  64*gw, 3,        2, node=18)(y18) #   19 (40,40,128) P3_downsample
    y20 = layers.Concatenate(name='19.Concat')([y19,y15])      #!  20 (40,40,256) P3_downsample
    y21 = C3   (y20,128*gw, 128*gw, 1*gd, False, node=20)      #   21 (40,40,256) P4_out

    y22 = Conv (    128*gw, 128*gw, 3,        2, node=21)(y21) #!  22 (20,20,256) P4_downsample
    y23 = layers.Concatenate(name='22.Concat')([y22,y11])      #   23 (20,20,512) P4_downsample
    y24 = C3   (y23,256*gw, 256*gw, 1*gd, False, node=23)      #!  24 (20,20,512) P5_out
    return Detect([y18,y21,y24],nc=nc,imgsz=imgsz,anchors=anchors,strides=strides,ch=ch,node=24)
    # return Detect([y24,y21,y18],nc=nc,imgsz=imgsz,anchors=anchors,strides=strides,ch=ch,node=24)

def yolov5(nc=80,imgsz=(640,640),strides=[8,16,32],anchors=None,ch=[128,256,512],gw=2,gd=1):
    x = keras.Input([imgsz[1],imgsz[0],3])
    feats = yolov5_backbone(x,gw=gw,gd=gd)
    ys = yolov5_head(
        *feats,
        nc=nc,
        imgsz=imgsz,
        strides=strides,
        anchors=anchors,
        ch=ch,
        gw=gw,
        gd=gd,
    )
    return keras.Model(inputs=x,outputs=ys)

def yolov5s(nc=80,imgsz=(640,640),strides=[8,16,32],anchors=None,ch=[128,256,512],gt=False):
    """
        imgsz(w,h)
    """
    if gt: nc += 1 #! 如果设置gt类别, 则自动增加默认类别, 总类别数加1
    return yolov5(nc=nc, imgsz=imgsz, strides=strides, anchors=anchors, ch=ch, gw=2, gd=1)

def yolov5m(nc=80,imgsz=(640,640),strides=[8,16,32],anchors=None,ch=[128,256,512],gt=False):
    """
        imgsz(w,h)
    """
    if gt: nc += 1 #! 如果设置gt类别, 则自动增加默认类别, 总类别数加1
    return yolov5(nc=nc, imgsz=imgsz, strides=strides, anchors=anchors, ch=ch, gw=3, gd=2)

def yolov5l(nc=80,imgsz=(640,640),strides=[8,16,32],anchors=None,ch=[128,256,512],gt=False):
    """
        imgsz(w,h)
    """
    if gt: nc += 1 #! 如果设置gt类别, 则自动增加默认类别, 总类别数加1
    return yolov5(nc=nc, imgsz=imgsz, strides=strides, anchors=anchors, ch=ch, gw=4, gd=3)

def yolov5x(nc=80,imgsz=(640,640),strides=[8,16,32],anchors=None,ch=[128,256,512],gt=False):
    """
        imgsz(w,h)
    """
    if gt: nc += 1 #! 如果设置gt类别, 则自动增加默认类别, 总类别数加1
    return yolov5(nc=nc, imgsz=imgsz, strides=strides, anchors=anchors, ch=ch, gw=5, gd=4)